import json
import os
from typing import Type, Any

import dotenv
import requests
from langchain.agents import Tool
from langchain_community.tools import GoogleSerperRun
from langchain_community.utilities import GoogleSerperAPIWrapper
from langchain_core.messages import ToolMessage
from langchain_core.prompts import ChatPromptTemplate
from pydantic import Field, BaseModel
from langchain_core.runnables import RunnablePassthrough
from langchain_core.tools import BaseTool
from langchain_openai import ChatOpenAI
from langchain.tools import tool

dotenv.load_dotenv()


## 博查工具，中国大陆内需要用来替换google搜索。 验证可用的
@tool
def bocha_web_search_tool(query: str, count: int = 8) -> str:
    """
    使用Bocha Web Search API进行联网搜索，返回搜索结果的字符串。

    参数:
    - query: 搜索关键词
    - count: 返回的搜索结果数量

    返回:
    - 搜索结果的字符串形式
    """
    url = 'https://api.bochaai.com/v1/web-search'
    headers = {
        'Authorization': f'Bearer {os.getenv("BOCHA_API_KEY")}',  # 请替换为你的API密钥
        'Content-Type': 'application/json'
    }
    data = {
        "query": query,
        "freshness": "noLimit",  # 搜索的时间范围，例如 "oneDay", "oneWeek", "oneMonth", "oneYear", "noLimit"
        "summary": True,  # 是否返回长文本摘要总结
        "count": count
    }

    response = requests.post(url, headers=headers, json=data)

    if response.status_code == 200:
        # 返回给大模型的格式化的搜索结果文本
        # 可以自己对博查的搜索结果进行自定义处理
        return str(response.json())
    else:
        raise Exception(f"API请求失败，状态码: {response.status_code}, 错误信息: {response.text}")


class GaodeWeatherArgsSchema(BaseModel):
    city: str = Field(description="需要查询天气预报的目标城市，例如：广州")


class GoogleSerperArgsSchema(BaseModel):
    query: str = Field(description="执行谷歌搜索的查询语句")


class GaodeWeatherTool(BaseTool):
    """根据传入的城市名查询天气"""
    name: str = "gaode_weather"
    description: str = "当你想查询天气或者天气相关的问题时可以使用的工具"
    args_schema: Type[BaseModel] = GaodeWeatherArgsSchema

    def _run(self, *args: Any, **kwargs: Any) -> str:
        """根据传入的城市名运行调用api获取城市对应的天气预报信息"""
        try:
            # 1. 获取高德API秘钥， 如果没有创建的话。则抛出错误
            gaode_api_key = os.getenv("GAODE_API_KEY")
            if not gaode_api_key:
                return f"高德开放平台api未获取"

            # 2. 从参数中获取city城市名字
            city = kwargs.get("city", "")
            api_domain = "https://restapi.amap.com/v3/"
            session = requests.session()

            # 3. 发起行政区域编码查询。根据city获取ad_code
            city_response = session.request(
                method="GET",
                url=f"{api_domain}config/district?key={gaode_api_key}&keywords={city}&subdistrict=1&extensions=base",
                headers={"Content-Type": "application/json;charset=uft-8"},
            )
            city_response.raise_for_status()
            city_data = city_response.json()
            if city_data.get("info") == "OK":
                ad_code = city_data["districts"][0]["adcode"]

                # 4. 根据得到的ad_code调用天气预报API接口，获取天气信息
                weather_response = session.request(
                    method="GET",
                    url=f"{api_domain}weather/weatherInfo?key={gaode_api_key}&city={ad_code}&extensions=all&output=json",
                    headers={"Content-Type": "application/json;charset=uft-8"},
                )
                weather_response.raise_for_status()
                weather_data = weather_response.json()
                if weather_data.get("info") == "OK":
                    # 5. 返回最后的结果字符串
                    return json.dumps(weather_data)
            return f"获取{city}天气预报信息失败"
        except Exception as e:
            return f"获取{kwargs.get('city', '')}天气预报信息失败"


# 1. 定义工具列表
### ① 高德天气工具
gaode_weather = GaodeWeatherTool()

### ② 谷歌查询工具 -- 需要VPN的
google_serper = GoogleSerperRun(
    name="google_serper",
    description=(
        "一个低成本的谷歌搜索API"
        "当你需要回答有关时事的问题时，可以调用该工具"
        "该工具的输入是搜索查询语句"
    ),
    args_schema=GoogleSerperArgsSchema,
    api_wrapper=GoogleSerperAPIWrapper(),
)

# ③ 博查国内的查询工具，替代谷歌，but不用vpn，就是麻烦点
bocha_tool = Tool(
    name="BochaWebSearch",
    func=bocha_web_search_tool,
    description="使用Bocha Web Search API进行网络搜索",
    args_schema=GoogleSerperArgsSchema,
    api_wrapper=GoogleSerperAPIWrapper(),
)

tool_dict = {
    gaode_weather.name: gaode_weather,
    # google_serper.name: google_serper,
    bocha_tool.name: bocha_tool,
}

tools = [tool for tool in tool_dict.values()]

# 2. 创建Prompt
prompt = ChatPromptTemplate.from_messages([
    ("system", "你是由OpenAI开发的聊天机器人，可以帮助用户回答问题，必要时可请调用工具帮助用户解答"),
    ("human", "{query}"),
])

# 3. 创建大预言模型并绑定工具
llm = ChatOpenAI(model="kimi-k2-0711-preview", temperature=0)
llm_tools = llm.bind_tools(tools=tools)

# 4. 创建链应用
chain = {"query": RunnablePassthrough()} | prompt | llm_tools

# 5. 调用链应用，并获取输出响应
query = "今年的阅兵当天北京天气如何，新装备都有哪些亮点？"
resp = chain.invoke(query)
tool_calls = resp.tool_calls

# 6.判断是工具调用还是正常输出结果
if len(tool_calls) <= 0:
    print("生成内容: ", resp.content)
else:
    # 7.将历史的系统消息、人类消息、AI消息组合
    messages = prompt.invoke(query).to_messages()
    messages.append(resp)
    print(messages)
    # 8.循环遍历所有工具调用信息
    for tool_call in tool_calls:
        tool = tool_dict.get(tool_call.get("name"))  # 获取需要执行的工具
        print("正在执行工具: ", tool.name)
        content = tool.invoke(tool_call.get("args"))  # 工具执行的内容/结果
        print("工具返回结果: ", content)
        tool_call_id = tool_call.get("id")
        messages.append(ToolMessage(
            content=content,
            tool_call_id=tool_call_id,
        ))
    print("输出内容: ", llm.invoke(messages).content)
